import os
import argparse
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from dataset import ECSSDDataset
from picanet_resnet18 import PiCANetResnet18
from train import train
from test import test
from metrics import evaluate_dataset
from evaluate import evaluate
import glob

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", type=str, choices=["train", "eval"], default="train", help="train or eval mode")
    parser.add_argument("--data_path", type=str, default="../ECSSD", help="path to ECSSD dataset")
    parser.add_argument("--save_path", type=str, default="./checkpoints", help="model save directory")
    parser.add_argument("--result_dir", type=str, default="./results", help="directory to save predicted masks")
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight", type=str, default="", help="path to model weight for evaluation")
    args = parser.parse_args()

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
    mask_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])

    image_dir = os.path.join(args.data_path, 'images')
    mask_dir = os.path.join(args.data_path, 'ground_truth_mask')
    dataset = ECSSDDataset(image_dir, mask_dir, transform, mask_transform)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = PiCANetResnet18().to(device)


    if args.mode == "train":
        train(args,device)
        weight_list = sorted([f for f in os.listdir(args.save_path) if f.endswith('.pth')])
        if not weight_list:
            print("No model found in save_path.")
            return
        model.load_state_dict(torch.load(os.path.join(args.save_path, weight_list[-1])))
    else:  # eval mode
        if args.weight == "":
            print("Please provide --weight path for evaluation.")
            return
        model.load_state_dict(torch.load(args.weight))

    test(model, dataloader, device,result_dir=args.result_dir)
    evaluate(args.result_dir, mask_dir)

if __name__ == "__main__":
    main()
